#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Hugging Face配置模块
确保在项目启动时正确设置环境变量，解决国内网络问题
"""

import os
import sys
import platform

def setup_hf_environment():
    """设置Hugging Face环境变量"""
    print("🔧 设置Hugging Face环境变量...")
    
    # 基础配置
    config = {
        # 镜像源配置
        'HF_ENDPOINT': 'https://hf-mirror.com',
        'HF_HUB_URL': 'https://hf-mirror.com',
        
        # 缓存目录配置
        'HF_HOME': os.path.expanduser('~/.cache/huggingface'),
        'TRANSFORMERS_CACHE': os.path.expanduser('~/.cache/huggingface/transformers'),
        'HF_DATASETS_CACHE': os.path.expanduser('~/.cache/huggingface/datasets'),
        
        # 禁用遥测和隐式token
        'HF_HUB_DISABLE_TELEMETRY': '1',
        'HF_HUB_DISABLE_IMPLICIT_TOKEN': '1',
        
        # 其他优化配置
        'HF_HUB_OFFLINE': '0',  # 允许在线下载
        'HF_HUB_DISABLE_SYMLINKS_WARNING': '1',  # 禁用符号链接警告
        'HF_HUB_DISABLE_PROGRESS_BARS': '0',  # 显示进度条
    }
    
    # 设置环境变量
    for key, value in config.items():
        os.environ[key] = value
        print(f"   {key} = {value}")
    
    # 创建缓存目录
    cache_dirs = [
        config['HF_HOME'],
        config['TRANSFORMERS_CACHE'],
        config['HF_DATASETS_CACHE']
    ]
    
    for cache_dir in cache_dirs:
        os.makedirs(cache_dir, exist_ok=True)
        print(f"   📁 创建目录: {cache_dir}")
    
    print("✅ Hugging Face环境变量设置完成")
    return config

def test_hf_connection():
    """测试Hugging Face连接"""
    print("🔍 测试Hugging Face连接...")
    
    try:
        import requests
        
        # 测试镜像源连接
        test_urls = [
            "https://hf-mirror.com/api/models",
            "https://hf-mirror.com/api/spaces"
        ]
        
        for url in test_urls:
            try:
                response = requests.get(url, timeout=10)
                if response.status_code == 200:
                    print(f"✅ 连接成功: {url}")
                    return True
                else:
                    print(f"❌ 连接失败: {url}, 状态码: {response.status_code}")
            except Exception as e:
                print(f"❌ 连接异常: {url}, 错误: {e}")
        
        return False
        
    except ImportError:
        print("❌ requests库未安装，跳过连接测试")
        return False

def get_hf_config():
    """获取当前Hugging Face配置"""
    config = {}
    hf_vars = [
        'HF_ENDPOINT', 'HF_HUB_URL', 'HF_HOME', 
        'TRANSFORMERS_CACHE', 'HF_DATASETS_CACHE',
        'HF_HUB_DISABLE_TELEMETRY', 'HF_HUB_DISABLE_IMPLICIT_TOKEN'
    ]
    
    for var in hf_vars:
        config[var] = os.environ.get(var, '未设置')
    
    return config

def print_hf_status():
    """打印Hugging Face状态信息"""
    print("📋 Hugging Face配置状态:")
    config = get_hf_config()
    
    for key, value in config.items():
        status = "✅" if value != "未设置" else "❌"
        print(f"   {status} {key}: {value}")
    
    # 测试连接
    connection_ok = test_hf_connection()
    print(f"   {'✅' if connection_ok else '❌'} 连接测试: {'成功' if connection_ok else '失败'}")

def setup_for_sentence_transformers():
    """为sentence-transformers设置特殊配置"""
    print("🔧 为sentence-transformers设置特殊配置...")
    
    # 设置sentence-transformers特定的环境变量
    os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.path.expanduser('~/.cache/sentence-transformers')
    
    # 创建sentence-transformers缓存目录
    st_cache_dir = os.environ['SENTENCE_TRANSFORMERS_HOME']
    os.makedirs(st_cache_dir, exist_ok=True)
    print(f"   📁 创建sentence-transformers缓存目录: {st_cache_dir}")
    
    # 设置模型下载参数
    os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '300'  # 5分钟超时
    os.environ['HF_HUB_DOWNLOAD_RETRY_DELAY'] = '1'  # 重试延迟1秒
    
    print("✅ sentence-transformers配置完成")

def create_hf_config_file():
    """创建Hugging Face配置文件"""
    config_content = """# Hugging Face配置文件
# 此文件包含Hugging Face相关的环境变量配置

# 镜像源配置
HF_ENDPOINT=https://hf-mirror.com
HF_HUB_URL=https://hf-mirror.com

# 缓存目录配置
HF_HOME=~/.cache/huggingface
TRANSFORMERS_CACHE=~/.cache/huggingface/transformers
HF_DATASETS_CACHE=~/.cache/huggingface/datasets
SENTENCE_TRANSFORMERS_HOME=~/.cache/sentence-transformers

# 功能配置
HF_HUB_DISABLE_TELEMETRY=1
HF_HUB_DISABLE_IMPLICIT_TOKEN=1
HF_HUB_DISABLE_SYMLINKS_WARNING=1
HF_HUB_OFFLINE=0

# 下载配置
HF_HUB_DOWNLOAD_TIMEOUT=300
HF_HUB_DOWNLOAD_RETRY_DELAY=1

# 模型优先级列表
MODEL_PRIORITY=[
    "shibing624/text2vec-base-chinese",
    "GanymedeNil/text2vec-large-chinese", 
    "paraphrase-multilingual-MiniLM-L12-v2",
    "all-MiniLM-L6-v2"
]
"""
    
    config_file = "hf_config.env"
    with open(config_file, 'w', encoding='utf-8') as f:
        f.write(config_content)
    
    print(f"✅ 配置文件创建成功: {config_file}")
    return config_file

def main():
    """主函数 - 用于独立运行配置"""
    print("🚀 Hugging Face配置工具")
    print("=" * 50)
    
    # 1. 设置环境变量
    config = setup_hf_environment()
    
    # 2. 设置sentence-transformers特殊配置
    setup_for_sentence_transformers()
    
    # 3. 创建配置文件
    config_file = create_hf_config_file()
    
    # 4. 打印状态
    print_hf_status()
    
    print("\n✅ 配置完成！")
    print("💡 现在可以正常使用sentence-transformers了")

# 当作为模块导入时，自动设置环境变量
if __name__ != "__main__":
    # 自动设置环境变量
    setup_hf_environment()
    setup_for_sentence_transformers()

if __name__ == "__main__":
    main() 